import { randomUUID } from "crypto";
import { ModelProviderFactory } from "../../model";
import { Logger } from "../../utils/Logger";
import { MessageHandler, MessageType } from "../../utils/MessageHandler";
import { AgentViewController } from "../AgentViewController";
import { GlobalStore } from "../../store/global";
import { wrapMessages } from "../../model/helper";
import { LocalChunkMessage, LocalModelMessage } from "../../store/types";
import { ChunkMessage } from "../../model/types";

/**
 * 消息处理器
 * 用于处理来自大模型的消息
 */
export class LLMMessageHandler implements MessageHandler {
	constructor(private readonly controller: AgentViewController) {}

	syncTask(taskId: string) {
		const globalStore = new GlobalStore(this.controller.context);
		this.controller.sendMessageToWebView({type: MessageType.getTask, data: globalStore.getTask(taskId)});	
	}

	async createMessage(message: any) {
		const context = this.controller.context;
		const globalStore = new GlobalStore(context);
		Logger.log("context: success");
		const msgList = message.data.message;
		Logger.log(JSON.stringify(msgList));

		// 读取本地消息
		const taskId = message.data.taskId;
		const localMessages = globalStore.getTask(taskId)?.messages?? [];
	
		const userMessage: LocalModelMessage = {
		  message: msgList[0],
			taskId: taskId,
			msgId: randomUUID(),
			createTime: Date.now(),
			updateTime: Date.now(),
		};
		localMessages?.push(userMessage);
		if(localMessages?.length <= 1) {
			Logger.log('AI Code Agent: no local messages');
			globalStore.saveTask({
				taskId: taskId,
				firstMessage: userMessage,
				messages: localMessages?? [],
			});
		} else {
			globalStore.saveTaskMessages(taskId, localMessages ?? []);
		}
		// 同步任务信息到webview
		this.syncTask(taskId);

		const msgId = randomUUID();
		let newChunks: ChunkMessage[] = [];
		const createTime = Date.now();
		const stream = await ModelProviderFactory.createMessage(msgList, context);
		if(stream != undefined) {
			let index = 0;
			for await (const chunk of stream) {
				Logger.log('AI Code Agent: received chunk from LLM');
				Logger.log(JSON.stringify(chunk));
				this.controller.sendMessageToWebView({type: MessageType.receivedLLMMessage, data: {
					taskId: message.data.taskId,
					msgId,
					chunk,
					index: index++,
				} as LocalChunkMessage});
				newChunks.push(chunk);
			}
		}

		newChunks = wrapMessages(newChunks) ?? [];
		const newMessage: LocalModelMessage = {
			message: {
				role: 'assistant',
				chunks: newChunks,
				content: newChunks.map((chunk) => chunk.content).join(''),
			},
			taskId: message.data.taskId,
			msgId,
			createTime: createTime,
			updateTime: Date.now(),
		};
		localMessages?.push(newMessage);
		globalStore.saveTaskMessages(taskId, localMessages ?? []);

		// 同步任务信息到webview
		this.syncTask(taskId);
	}

	handleMessage(message: any): void {
		Logger.log('AI Code Agent: handle LLM message');
		Logger.log(JSON.stringify(message));
		switch(message.type) {
			case MessageType.sendMessageToLLM:
				{
					this.createMessage(message);
					break;
				}
		}
	}
}